from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_openai import  OpenAIEmbeddings
from langchain_core.runnables import RunnableLambda
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 准备测试数据，假设我们提供的文档数据如下:
documents = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好而闻名。", # 文档内容
        metadata={"source":"哺乳动物宠物文档"} # 文档元数据
    ),
    Document(
        page_content="猫是独立的宠物，通常喜欢自己的空间。",
        metadata={"source":"哺乳动物宠物文档"}
    ),
    Document(
        page_content="金鱼实时初学者的流行宠物，需要相对简单的护理。",
        metadata={"source": "鱼类宠物文档"}
    ),
    Document(
        page_content="白头鹰是犀利的猎手，视力很强大。",
        metadata={"source": "鸟类宠物文档"}
    )
]

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200, add_start_index=True
)
# 对文档进行分块
all_splits = text_splitter.split_documents(documents)

# 实例化向量数据库
vector_db = Chroma.from_documents(all_splits,embedding=OpenAIEmbeddings(model="text-embedding-3-large"))

# 相似度查询,返回相似的分数，分数越低，相似度越高
# print(vector_db.similarity_search_with_score("加菲猫"))

# 相似度查询,返回相似的分数，分数越低，相似度越高 异步方式
result =  vector_db.asimilarity_search("加菲猫")
print(result)

# 相似度查询,返回相似的文档
retriever = RunnableLambda(vector_db.similarity_search).bind(k=1)
print(retriever.batch(['加菲猫', '鲸鱼']))
